# -*- coding:utf-8 -*-

# @Time    : 2018/11/19 2:31 PM

# @Author  : Swing


import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

from keras.layers.core import Dense, Flatten
from keras.layers.convolutional import Conv2D
from keras.layers.pooling import MaxPool2D

from keras.objectives import categorical_crossentropy

import keras.backend as K

data_dir = '/Users/zhubin/Documents/ai/data/mnist/'
mnist = input_data.read_data_sets(data_dir, one_hot=True)

x = tf.placeholder(tf.float32, [None, 784], name='x')
y_ = tf.placeholder(tf.float32, [None, 10])

learning_rate = tf.placeholder(tf.float32)

with tf.name_scope('reshape'):
    x_image = tf.reshape(x, [-1, 28, 28, 1])

net = Conv2D(32, [5, 5], strides=[1, 1], activation='relu', padding='SAME', input_shape=[28, 28, 1], name='conv1')(x_image)

net = MaxPool2D(pool_size=[2, 2])(net)

net = Conv2D(64, [5, 5], strides=[1, 1], activation='relu', padding='SAME')(net)

net = MaxPool2D(pool_size=[2, 2])(net)

net = Flatten()(net)

net = Dense(1000, activation='relu')(net)

net = Dense(10, activation='softmax')(net)

cross_entropy = tf.reduce_mean(categorical_crossentropy(y_, net))
l2_loss = tf.add_n([tf.nn.l2_loss(w) for w in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)])

total_loss = cross_entropy + 7e-5 * l2_loss

sess = tf.Session()

K.set_session(sess)

train_step = tf.train.GradientDescentOptimizer(learning_rate).minimize(total_loss)

init_op = tf.global_variables_initializer()

sess.run(init_op)

# Train
for step in range(3000):
    batch_xs, batch_ys = mnist.train.next_batch(100)
    lr = 0.01
    _, loss, l2_loss_value, total_loss_value = sess.run(
        [train_step, cross_entropy, l2_loss, total_loss],
        feed_dict={x: batch_xs, y_: batch_ys, learning_rate: lr})

    if (step + 1) % 100 == 0:
        print('step %d, entropy loss: %f, l2_loss: %f, total loss: %f' %
              (step + 1, loss, l2_loss_value, total_loss_value))
        # Test trained model
        correct_prediction = tf.equal(tf.argmax(net, 1), tf.argmax(y_, 1))
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
        print(sess.run(accuracy, feed_dict={x: batch_xs, y_: batch_ys}))
    if (step + 1) % 1000 == 0:
        print(sess.run(accuracy, feed_dict={x: mnist.test.images,
                                            y_: mnist.test.labels}))
